import cv2
import numpy as np
import matplotlib.pyplot as plt

def canny_edge_detector(img,hi_threshold,low_threshold):
    EDGE_VAL=hi_threshold+1
    """step 1 高斯滤波"""
    sigma=3;
    img = cv2.GaussianBlur(img,(sigma,sigma),0)

    """step 2.1  梯度幅值"""
    dx = cv2.Sobel(img,cv2.CV_16S,1,0)      #向右为x轴正方向 ,None,3,1,0,cv2.BORDER_REPLICATE
    dy = cv2.Sobel(img,cv2.CV_16S,0,1)      #向下为y轴正方向

    #mag=np.hypot(dx,dy)   #计算L2范数，即梯度的幅值
    mag=np.abs(dx)+np.abs(dy)

    """step 2.2 & 3 梯度方向计算， 非最大抑制"""
    TG22=0.4142135623731
    TG67=2.4142135623731
    nHeight, nWidth = img.shape[:2]
    edge=np.zeros([nHeight,nWidth])
    tangent_test=np.zeros([nHeight,nWidth])
    for r in range(1,nHeight-1):    #rows,y轴
        for c in range(1,nWidth-1): #columns,x轴
            p1=[0,0]; p2=[0,0]
            #计算正切值
            if(dx[r,c]==0) :    tangent=5;                   tangent_test[r,c]=5
            else:               tangent=dy[r,c]/dx[r,c];     tangent_test[r,c]=tangent
            #将方向舍入到0，90，45，135四个方向
            if np.abs(tangent)<TG22 :   #水平, 取东西上邻近像素，进行非最大抑制比较
                #if( mag[r,c-1] <mag[r,c] and mag[r,c+1] <=mag[r,c] ):   #判断中，必须有一个是小于等于，防止边缘幅值处都相等的情况
                    #edge[r,c]=mag[r,c]
                y1=r; x1=c+1; y2=r; x2=c-1
            elif TG22<=tangent and tangent<=TG67:   #45度(注意y轴朝下)
                #if mag[r-1,c-1] <mag[r,c] and mag[r+1,c+1]<=mag[r,c]:
                    #edge[r,c]=mag[r,c]
                y1=r+1; x1=c+1; y2=r-1; x2=c-1
            elif -TG67<=tangent and tangent<=-TG22: #135度
                #if mag[r+1, c-1] < mag[r, c] and mag[r-1,c+1] <= mag[r, c]:
                    #edge[r,c]=mag[r,c]
                y1=r-1;x1=c+1; y2=r+1;x2=c-1
            else:     #垂直，取南北邻近像素
                #if(mag[r-1,c] <mag[r,c] and mag[r+1,c]<=mag[r,c]):
                    #edge[r,c]=mag[r,c]
                y1=r+1;x1=c; y2=r-1;x2=c
            if mag[r,c]>mag[y2,x2] and ( mag[r,c]>mag[y1,x1] or ( mag[r,c]==mag[y1,x1] and (( x1==c+1 and y1==r) or (x1==c and y1==r+1)) ) ):
                edge[r,c]=mag[r,c]

    plt.gray()
    #plt.subplot(221)
    plt.imshow(mag,interpolation="nearest")
    plt.show()
    #plt.subplot(222)
    plt.imshow(edge,interpolation="nearest")
    #plt.subplot(223)
    plt.show()
    plt.imshow(tangent_test,interpolation="nearest")
    plt.show()

    """step 4 双阀值处理，生成强边"""
    strong_pxs=[]
    for r in range(1,nHeight-1):
        for c in range(1,nWidth-1):
            if edge[r,c] >hi_threshold:
                strong_pxs.append([r,c])
                edge[r,c]=EDGE_VAL

    """step 5 弱边处理，通过强边像素追踪弱边像素，形成直实边缘"""
    while strong_pxs:
        r,c=strong_pxs.pop(0)
        neighbor={(-1,-1),(-1,0),(-1,1),(1,-1),(1,0),(1,1),(0,-1),(0,1)}
        #查看其邻域内有没有弱边像素，将弱边像素也标记为强边像素
        for rx,cx in neighbor:
            if edge[r+rx,c+cx]>low_threshold and edge[r+rx,c+cx]<hi_threshold:
                edge[r+rx,c+cx]=EDGE_VAL;
                strong_pxs.append([r+rx,c+cx])
    ##剔除未与强边像素连通的弱边像素
    for r in range(1,nHeight-1):
        for c in range(1,nWidth-1):
            if edge[r,c]<=hi_threshold: edge[r,c]=0;

    #变成0，1
    edge=cv2.convertScaleAbs(edge,None,1/EDGE_VAL)

    return edge



"""
##参数
fileName=r'misc_pic\yl01_01.jpg'
#fileName=r'misc_pic\empire.jpg'
high_threshold =70
low_threshold=30

##载入图像
img = cv2.imread(fileName)
img= cv2.cvtColor(img,cv2.COLOR_RGB2GRAY)

plt.gray()
plt.imshow(img,interpolation="nearest")
plt.show()

edge=canny_edge_detector(img,high_threshold,low_threshold)

plt.gray()
plt.imshow(edge,interpolation="nearest")
plt.show()
"""